#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import json
import os
import sys
import asyncio
from datetime import datetime
from typing import Dict, List, Any, Optional
from xmlrpc.client import boolean

from pydantic import BaseModel, Field

from src.utils.env_config import get_task_name, file_state, file_context

"""
任务1：固定任务拆解与状态管理 (使用AgentScope框架)
实现报告生成流程的固定子任务拆解、执行状态记录、结构化报告框架输出
"""

class StateManager:
    """状态管理器：负责任务状态的持久化存储和管理"""

    def __init__(self, state_file: str = file_state, context_file: str = "context.json"):
        self.state_file = state_file
        self.state_data = self._load_state()
        self.context_file = context_file
        self.context_data = self._load_context()

    def _load_context(self) -> Dict[str, Any]:
        """加载上下文文件"""
        if os.path.exists(self.context_file):
            with open(self.context_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        return self._create_initial_context()

    @staticmethod
    def _create_initial_context() -> Dict[str, Any]:
        """创建初始上下文"""
        return {
            "context": {
                "topic": ""
            }
        }

    def _load_state(self) -> Dict[str, Any]:
        """加载状态文件"""
        if os.path.exists(self.state_file):
            with open(self.state_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        return self._create_initial_state()

    @staticmethod
    def _create_initial_state() -> Dict[str, Any]:
        """创建初始状态"""
        return {"task_list": []}
    def clear_task_status(self, task_id: str):
        """清除任务状态"""
        for task in self.state_data["task_list"]:
            if task["task_id"] == task_id:
                task["status_history"] = []
                break
    def clear_context(self):
        """清除所有上下文数据"""
        self.context_data = self._create_initial_context()
        self.save_to_context_data()

    def update_task_status(self, task_id: str, status: str) -> None:
        """更新任务状态"""
        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        change =  False
        for task in self.state_data["task_list"]:
            if task["task_id"] == task_id:
                task["status_history"].append({
                    "status": status,
                    "timestamp": current_time
                })
                change = True
                break
        ## 添加任务记录
        if not change:
            task_name = get_task_name(task_id)
            self.state_data["task_list"].append({
                "task_id": task_id,
                "task_name": task_name,
                "status_history": [
                    {"status": status, "timestamp": current_time}
                ]
            })
        self._save_state()

    def get_task_status(self, task_id: str) -> str:
        """获取任务状态"""
        for task in self.state_data["task_list"]:
            if task["task_id"] == task_id:
                status_history = task["status_history"]
                if status_history:
                    return status_history[-1]["status"]
        return ""

    def save_context(self, key: str, value: Any) -> None:
        """保存上下文数据"""
        self.context_data["context"][key] = value
        """保存上下文到文件"""
        with open(self.context_file, 'w', encoding='utf-8') as f:
            json.dump(self.context_data, f, ensure_ascii=False, indent=2)
        f.close()
        self.save_to_context_data()


    def get_context(self, key: str) -> Any:
        """获取上下文数据"""
        return self.context_data["context"].get(key)

    def _save_state(self) -> None:
        """保存状态到文件"""
        with open(self.state_file, 'w', encoding='utf-8') as f:
            json.dump(self.state_data, f, ensure_ascii=False, indent=2)
        f.close()

    def save_to_context_data(self) -> None:
        """保存数据到上下文数据"""
        new_context = {}
        if self.context_data["context"].get("keywords"):
            new_context["T1_output"] = {"result":self.context_data["context"]["keywords"]}

        if self.context_data["context"].get("retrieval_summary"):
            new_context["T2_output"] = {"result":self.context_data["context"]["retrieval_summary"]}

        if self.context_data["context"].get("case_analysis"):
            new_context["T5_output"] = {"result":self.context_data["context"]["case_analysis"]}

        if self.context_data["context"].get("policy_impact"):
            new_context["T6_output"] = {"result":self.context_data["context"]["policy_impact"]}

        if self.context_data["context"].get("abstract"):
            new_context["T3_output"] = {"result": self.context_data["context"]["abstract"]}

        if self.context_data["context"].get("markdown_content"):
            new_context["T4_output"] = {"result":self.context_data["context"]["markdown_content"]}

        """保存上下文到文件"""
        with open(file_context, 'w', encoding='utf-8') as f:
            json.dump(new_context, f, ensure_ascii=False, indent=2)

        f.close()